#   kmeans虽然会对数据集中的框进行聚类，但是很多数据集由于框的大小相近，聚类出来的9个框相差不大，
#   这样的框反而不利于模型的训练。因为不同的特征层适合不同大小的先验框，越浅的特征层适合越大的先验框
#   原始网络的先验框已经按大中小比例分配好了，不进行聚类也会有非常好的效果。
import glob
import xml.etree.ElementTree as ET
import config as cfg
import numpy as np

def cas_iou(box,cluster):
    x = np.minimum(cluster[:,0],box[0])
    y = np.minimum(cluster[:,1],box[1])

    intersection = x * y
    area1 = box[0] * box[1]

    area2 = cluster[:,0] * cluster[:,1]
    iou = intersection / (area1 + area2 -intersection)

    return iou

def avg_iou(box,cluster):
    return np.mean([np.max(cas_iou(box[i],cluster)) for i in range(box.shape[0])])

def kmeans(box,k):
    #-------------------------------------------------------------#
    #   取出一共有多少框
    #-------------------------------------------------------------#
    row = box.shape[0]
    
    #-------------------------------------------------------------#
    #   每个框各个点的位置
    #-------------------------------------------------------------#
    distance = np.empty((row,k))
    
    #-------------------------------------------------------------#
    #   最后的聚类位置
    #-------------------------------------------------------------#
    last_clu = np.zeros((row,))

    np.random.seed()

    #-------------------------------------------------------------#
    #   随机选5个当聚类中心
    #-------------------------------------------------------------#
    cluster = box[np.random.choice(row,k,replace = False)]
    while True:
        #-------------------------------------------------------------#
        #   计算每一行距离五个点的iou情况。
        #-------------------------------------------------------------#
        for i in range(row):
            distance[i] = 1 - cas_iou(box[i],cluster)
        
        #-------------------------------------------------------------#
        #   取出最小点
        #-------------------------------------------------------------#
        near = np.argmin(distance,axis=1)

        if (last_clu == near).all():
            break
        
        #-------------------------------------------------------------#
        #   求每一个类的中位点
        #-------------------------------------------------------------#
        for j in range(k):
            cluster[j] = np.median(
                box[near == j],axis=0)

        last_clu = near

    return cluster

def load_data(path):
    '''
    路径为xml的父目录

    return:ndarray
        返回所有xml的所有object的box的宽高
    '''
    data = []
    #-------------------------------------------------------------#
    #   对于每一个xml都寻找box
    #-------------------------------------------------------------#
    for xml_file in glob.glob('{}/*xml'.format(path)):
        tree = ET.parse(xml_file)
        height = int(tree.findtext('./size/height'))
        width = int(tree.findtext('./size/width'))
        if height<=0 or width<=0:
            continue
        
        #-------------------------------------------------------------#
        #   对于每一个目标都获得它的宽高
        #-------------------------------------------------------------#
        for obj in tree.iter('object'):
            xmin = int(float(obj.findtext('bndbox/xmin'))) / width
            ymin = int(float(obj.findtext('bndbox/ymin'))) / height
            xmax = int(float(obj.findtext('bndbox/xmax'))) / width
            ymax = int(float(obj.findtext('bndbox/ymax'))) / height

            xmin = np.float64(xmin)
            ymin = np.float64(ymin)
            xmax = np.float64(xmax)
            ymax = np.float64(ymax)
            # 得到宽高
            data.append([xmax-xmin,ymax-ymin])
    return np.array(data)

def load_json_data(img_path,json_path):
    '''
    该函数用来读取一个bdd的json文件对象，获取其对应的每个图片对象的含有目标的bounding boxes的宽和高，
    且该宽高会被归一化到想对于原始图像宽高
    img_path:str
        图片的目录路径
    json_path:str
        json文件路径
    return:ndarray
        返回的一个ndarray数组，包含了宽高的信息
    '''
    import json
    from tqdm import tqdm
    import os
    import cv2
    data = []
    with open(json_path) as f:
        json_file = json.load(f)
        for img_info in tqdm(json_file):
            img_name = img_info['name']  # 图片对应的文件名
            # ---------------------读取完毕-----------------------------------
            img_full_path = os.path.join(img_path, img_name)
            img = cv2.imread(img_full_path)
            img_h,img_w = img.shape[:2]

            for img_label in img_info['labels']:
                if 'box2d' in img_label:
                    x1, y1 = img_label['box2d']['x1'], img_label['box2d']['y1']
                    x2, y2 = img_label['box2d']['x2'], img_label['box2d']['y2']
                    bbox_w = np.float64((x2-x1)/img_w)
                    bbox_h = np.float64((y2-y1)/img_h)
                    data.append([bbox_w,bbox_h])
    return np.array(data)



if __name__ == '__main__':
    img_path = '/media/jiji/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/private个人数据集/BDD10K/images'
    json_path = r'/media/jiji/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/private个人数据集/BDD10K/labels/bdd10k_json.json'

    #   运行该程序会生成yolo_anchors.txt
    SIZE        = cfg.input_img_size[0]
    anchors_num = 9

    #   存储格式为转化为比例后的width,height
    # data = load_data(path)    #如果加载一个xml文件夹，使用这个函数
    data = load_json_data(img_path,json_path)   #如果加载类似于bdd100k的一个json文件，使用这个函数

    #   使用k聚类算法
    out = kmeans(data,anchors_num)
    out = out[np.argsort(out[:,0])]
    print('acc:{:.2f}%'.format(avg_iou(data,out) * 100))
    print(out*SIZE)
    data = out*SIZE
    f = open("model_data/yolo_bdd_anchors.txt", 'w')
    row = np.shape(data)[0]
    for i in range(row):
        if i == 0:
            x_y = "%d,%d" % (data[i][0], data[i][1])
        else:
            x_y = ", %d,%d" % (data[i][0], data[i][1])
        f.write(x_y)
    f.close()
